#include <sys/socket.h>
#include <netinet/in.h>
#include <stdlib.h>
#include <poll.h>
#include <stdio.h>
#include <unistd.h>
#include <stdint.h>
#include <string.h>
#include <vector>
#include <string>

#define PORT 12346

struct Client {
    uint8_t size_size;
    uint32_t size;
    uint32_t msg_pos;
    char* msg;
    uint32_t out_size;
    uint32_t out_pos;
    char* out_msg;
    std::string name;
    Client(): size_size(0), msg_pos(0), msg(0), out_size(0), out_pos(0), out_msg(0) {}
};

struct Message {
    char* msg;
    uint32_t len;
    Message(char* _msg, uint32_t _len): msg(_msg), len(_len) {}
};

std::vector<Client> clients;
std::vector<pollfd> fds;
std::vector<Message> messages;

void shutdown() {
    for (int i = 0; i < fds.size(); ++i) {
        close(fds[i].fd);
    }
    exit(1);
}

void end_session(int i) {
    close(fds[i].fd);
    delete [] clients[i - 2].msg;
    delete [] clients[i - 2].out_msg;
    if (i != fds.size() - 1) {
        fds[i] = fds[fds.size() - 1];
        clients[i - 2] = clients[clients.size() - 1];
    }
    fds.pop_back();
    clients.pop_back();
}

void getRequest(int i, uint32_t from) {
    uint32_t size = 0;
    Client& c = clients[i - 2];
    for (int j = from; j < messages.size(); ++j) {
        size += messages[j].len;
    }
    c.out_msg = new char[size + 15];
    int s = sprintf(c.out_msg + 4, "%d", messages.size());
    *(uint32_t*)c.out_msg = size + s + 1;
    s += 4;
    c.out_msg[s++] = 0;
    for (int j = from; j < messages.size(); ++j) {
        strcpy(c.out_msg + s, messages[j].msg);
        s += messages[j].len;
    }
    c.out_size = s;
    fds[i].events = POLLIN | POLLOUT;
}

void process_msg(int i) {
    Client& c = clients[i - 2];
    c.msg[c.size] = 0;
    if (strncmp(c.msg, "bye", 3) == 0) {
        end_session(i);
        return;
    }
    if (strncmp(c.msg, "send", 4) == 0) {
        char* msg = new char[c.size + c.name.length()];
        msg[0] = '<';
        strcpy(msg + 1, c.name.c_str());
        msg[1 + c.name.length()] = '>';
        msg[2 + c.name.length()] = ' ';
        strncpy(msg + 3 + c.name.length(), c.msg + 4, c.size - 4);
        msg[c.size + c.name.length() - 1] = 0;
        messages.push_back(Message(msg, c.size + c.name.length()));
    } else
    if (!c.out_msg)
    if (strncmp(c.msg, "hello", 5) == 0) {
        if (c.msg[c.size - 1] != '1') {
            c.out_msg = new char[5];
            *(uint32_t*)c.out_msg = 1;
            c.out_msg[4] = 'n';
            c.out_size = 5;
            fds[i].events = POLLOUT;
        } else {
            c.name = std::string(c.msg + 5, c.size - 6);
            if (c.name.empty()) {
                c.name = "anonymous";
            }
            c.out_msg = new char[30];
            c.out_msg[4] = 'y';
            c.out_msg[5] = 0;
            int s = sprintf(c.out_msg + 6, "%d", messages.size());
            c.out_msg[6 + s] = 0;
            strcpy(c.out_msg + 7 + s, "list,history");
            c.out_size = 19 + s;
            *(uint32_t*)c.out_msg = c.out_size - 4;
            fds[i].events = POLLIN | POLLOUT;
        }
    } else
    if (strncmp(c.msg, "get", 3) == 0) {
        uint32_t n = 0;
        if (sscanf(c.msg + 3, "%u", &n) == 1) {
            getRequest(i, n);
        }
    } else
    if (strncmp(c.msg, "list", 4) == 0) {
        uint32_t pos = 9;
        c.out_size = pos;
        for (int j = 0; j < clients.size(); ++j) {
            c.out_size += clients[j].name.length() ? clients[j].name.length() + 1 : 0;
        }
        c.out_msg = new char[c.out_size];
        *(uint32_t*)c.out_msg = c.out_size - 4;
        strcpy(c.out_msg + 4, "list");
        c.out_msg[pos - 1] = 0;
        for (int j = 0; j < clients.size(); ++j) {
            if (!clients[j].name.empty()) {
                strcpy(c.out_msg + pos, clients[j].name.c_str());
                pos += clients[j].name.length();
                c.out_msg[pos++] = 0;
            }
        }
        fds[i].events = POLLIN | POLLOUT;
    } else
    if (strncmp(c.msg, "history", 7) == 0) {
        getRequest(i, 0);
    }
    c.size_size = 0;
    c.msg_pos = 0;
    delete [] c.msg;
    c.msg = 0;
}

void send_msg(int i) {
    Client& c = clients[i - 2];
    ssize_t r = send(fds[i].fd, c.out_msg + c.out_pos, c.out_size - c.out_pos, 0);
    if (r <= 0) {
        end_session(i);
        return;
    }
    c.out_pos += r;
    if (c.out_pos >= c.out_size) {
        if (c.out_size == 5 && c.out_msg[4] == 'n') {
            end_session(i);
        } else {
            c.out_pos = 0;
            delete [] c.out_msg;
            c.out_msg = 0;
            fds[i].events = POLLIN;
        }
    }
}

void recv_msg(int i) {
    Client& c = clients[i - 2];
    if (c.size_size < 4) {
        ssize_t r = recv(fds[i].fd, &c.size + c.size_size, 4 - c.size_size, 0);
        if (r <= 0) {
            end_session(i);
            return;
        }
        c.size_size += r;
        if (c.size_size == 4) {
            c.msg = new char[c.size + 1];
        }
        return;
    }
    ssize_t r = recv(fds[i].fd, c.msg, c.size - c.msg_pos, 0);
    if (r <= 0) {
        end_session(i);
        return;
    }
    c.msg_pos += r;
    if (c.msg_pos >= c.size) {
        process_msg(i);
    }
}

void loop() {
    int count = poll(fds.data(), fds.size(), -1);
    if (count < 0) {
        perror("poll");
        shutdown();
    }
    if (fds[0].revents & POLLERR || fds[0].revents & POLLHUP || fds[0].revents & POLLNVAL) {
        --count;
        fprintf(stderr, "Error on listener socket\n");
        shutdown();
    } else
    if (fds[0].revents & POLLIN) {
        --count;
        pollfd fd;
        fd.fd = accept(fds[0].fd, 0, 0);
        if (fd.fd >= 0) {
            fd.events = POLLIN;
            fds.push_back(fd);
            clients.push_back(Client());
        }
    }
    if (fds[1].revents & POLLERR || fds[1].revents & POLLHUP || fds[1].revents & POLLNVAL) {
        --count;
        fprintf(stderr, "Error on broadcast socket\n");
        fds[1].events = 0;
    } else
    if (fds[1].revents & POLLIN) {
        --count;
        struct sockaddr addr;
        socklen_t addrlen = sizeof(addr);
        char msg[32];
        int len = recvfrom(fds[1].fd, msg, sizeof(msg), 0, &addr, &addrlen);
        if (len > 0) {
            sendto(fds[1].fd, msg, len, 0, &addr, sizeof(sockaddr_in));
        }
    }
    for (int i = 2; count > 0; ++i) {
        if (fds[i].revents & POLLERR || fds[i].revents & POLLHUP || fds[i].revents & POLLNVAL) {
            --count;
            end_session(i--);
        } else if (fds[i].revents & POLLOUT || fds[i].revents & POLLIN) {
            --count;
            if (fds[i].revents & POLLOUT) send_msg(i);
            if (fds[i].revents & POLLIN) recv_msg(i);
        }
    }
}

int main() {
    int listen_sock = socket(AF_INET, SOCK_STREAM, 0);
    if (listen_sock < 0) {
        perror("socket");
        shutdown();
    }
    fds.push_back(pollfd());
    fds[0].fd = listen_sock;
    fds[0].events = POLLIN;
    
    int broadcast_sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (broadcast_sock < 0) {
        perror("socket");
        shutdown();
    }
    fds.push_back(pollfd());
    fds[1].fd = broadcast_sock;
    fds[1].events = POLLIN;
    
    struct sockaddr_in addr;
    addr.sin_family = AF_INET;
    addr.sin_port = htons(PORT);
    addr.sin_addr.s_addr = INADDR_ANY;
    if (bind(listen_sock, (struct sockaddr*)&addr, sizeof(addr)) < 0) {
        perror("bind");
        shutdown();
    }
    if (bind(broadcast_sock, (struct sockaddr*)&addr, sizeof(addr)) < 0) {
        perror("bind");
        shutdown();
    }  
    
    if (listen(listen_sock, 10) < 0) {
        perror("listen");
        shutdown();
    }
    
    while (1) loop();
}
